! ==============================================================================
! This file is part of CUPP (Community Unified Physics Package) since 2023.
! ==============================================================================

module cu_tdk_mod

  use cupp_mesh_mod
  use cu_const_mod
  use atm_types_mod

  implicit none

  private

  public cu_tdk_init
  public cu_tdk_final
  public cu_tdk_run

  type tdk_data_type
    integer :: id                            ! ID of the data
    type(cupp_mesh_type), pointer :: mesh    ! Dimensions of the data
    real(r8), allocatable :: env_t_lev (:,:) ! Environment temperature at half levels (K)
    real(r8), allocatable :: env_qs    (:,:) ! Environment saturated specific humidity (kg kg-1)
    real(r8), allocatable :: qte       (:,:) ! Moisture convergence (kg kg-1 s-1)
    integer , allocatable :: k_min_omg (:  ) ! Level index of minimum omega
    integer , allocatable :: k_cld_base(:  ) ! Level index of cloud base
    integer , allocatable :: k_cld_top (:  ) ! Level index of cloud top
    integer , allocatable :: level_type(:,:) ! Level type - 1: Sub-cloud level; 2: Condensation level (cloud base)
    integer , allocatable :: conv_type (:,:) ! Convection type - 1: Penetrative; 2: Stratocumulus; 3: Mid-level
    real(r8), allocatable :: cld_t     (:,:) ! Cloud temperature (K)
    real(r8), allocatable :: cld_qv    (:,:) ! Cloud specific humidity (kg kg-1)
    real(r8), allocatable :: cld_qc    (:,:) ! Cloud liquid water (kg kg-1)
    real(r8), allocatable :: cld_u     (:,:) ! Cloud u-wind (m s-1)
    real(r8), allocatable :: cld_v     (:,:) ! Cloud v-wind (m s-1)
    real(r8), allocatable :: mfub      (:  ) ! Updraft mass flux at cloud base (kg m-2 s-1)
    real(r8), allocatable :: mfu       (:,:) ! Updraft mass flux (kg m-2 s-1)
    real(r8), allocatable :: mfus      (:,:) ! Update flux of dry static energy
    real(r8), allocatable :: mfuq      (:,:) ! Update flux of specific humidity
    real(r8), allocatable :: mful      (:,:) ! Update flux of cloud liquid water
    real(r8), allocatable :: entr      (:  ) ! Entrainment rate (s-1)
    real(r8), allocatable :: detr_qc   (:,:) ! Cloud liquid water returned to environment by detrainment
    ! Tendencies
    real(r8), allocatable :: dptdt     (:,:) ! Potential temperature tendency (K s-1)
    real(r8), allocatable :: dqvdt     (:,:) ! Specific humidity tendency (kg kg-1 s-1)
    real(r8), allocatable :: dudt      (:,:) ! U-wind tendency (m s-2)
    real(r8), allocatable :: dvdt      (:,:) ! V-wind tendency (m s-2)
  contains
    procedure :: init => tdk_data_init
    procedure :: clear => tdk_data_clear
    final :: tdk_data_final
  end type tdk_data_type

  type(tdk_data_type), allocatable :: tdk_data(:)

contains

  subroutine cu_tdk_init(namelist, mesh)

    character(*), intent(in) :: namelist
    type(cupp_mesh_type), intent(in) :: mesh(:)

    integer iblk

    call cu_tdk_final()

    allocate(tdk_data(size(mesh)))
    do iblk = 1, size(mesh)
      call tdk_data(iblk)%init(mesh(iblk))
    end do

  end subroutine cu_tdk_init

  subroutine cu_tdk_final()

    if (allocated(tdk_data)) deallocate(tdk_data)

  end subroutine cu_tdk_final

  subroutine cu_tdk_run(state, tend, dt_phys, dt_cu)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_cu

    integer icol, ilev

    associate (mesh => state%mesh, tdk => tdk_data(state%id))
    do ilev = 2, mesh%nlev
      do icol = 1, mesh%ncol
        tdk%env_t_lev(icol,ilev) = ( &
          max(cpd * state%t(icol,ilev-1) + g * state%z(icol,ilev-1), &
              cpd * state%t(icol,ilev  ) + g * state%z(icol,ilev  )) &
          - g * state%z_lev(icol,ilev) &
        ) / cpd
      end do
    end do
    end associate

  end subroutine cu_tdk_run

  subroutine tdk_data_init(this, mesh)

    class(tdk_data_type), intent(inout) :: this
    type(cupp_mesh_type), intent(in), target :: mesh

    call this%clear()

    this%id   = mesh%id
    this%mesh => mesh

    allocate(this%env_t_lev(mesh%ncol,mesh%nlev))

  end subroutine tdk_data_init

  subroutine tdk_data_clear(this)

    class(tdk_data_type), intent(inout) :: this

    this%id   = 0
    this%mesh => null()

    if (allocated(this%env_t_lev)) deallocate(this%env_t_lev)

  end subroutine tdk_data_clear

  subroutine tdk_data_final(this)

    type(tdk_data_type), intent(inout) :: this

    call this%clear()

  end subroutine tdk_data_final

end module cu_tdk_mod